import re
import threading
import time

import numpy as np
import soundfile as sf
import yaml

from sanic import Sanic, response
from sanic.exceptions import ServerError
from sanic.log import logger as log
import argparse

app = Sanic("Example")

import tensorflow as tf

from tensorflow_tts.inference import AutoProcessor
from tensorflow_tts.inference import TFAutoModel
import sounddevice as sd

"""### Tacotron2"""
tacotron2 = TFAutoModel.from_pretrained("tensorspeech/tts-tacotron2-baker-ch", name="tacotron2")
"""### FastSpeech2"""
fastspeech2 = TFAutoModel.from_pretrained("tensorspeech/tts-fastspeech2-baker-ch", name="fastspeech2")
"""### MB-MelGAN"""
mb_melgan = TFAutoModel.from_pretrained("tensorspeech/tts-mb_melgan-baker-ch", name="mb_melgan")
"""## Inference"""
processor = AutoProcessor.from_pretrained("tensorspeech/tts-tacotron2-baker-ch")


def do_synthesis(input_text, text2mel_model, vocoder_model, text2mel_name, vocoder_name):
    input_ids = processor.text_to_sequence(input_text, inference=True)

    # text2mel part
    if text2mel_name == "TACOTRON":
        _, mel_outputs, stop_token_prediction, alignment_history = text2mel_model.inference(
            tf.expand_dims(tf.convert_to_tensor(input_ids, dtype=tf.int32), 0),
            tf.convert_to_tensor([len(input_ids)], tf.int32),
            tf.convert_to_tensor([0], dtype=tf.int32)
        )
    elif text2mel_name == "FASTSPEECH2":
        mel_before, mel_outputs, duration_outputs, _, _ = text2mel_model.inference(
            tf.expand_dims(tf.convert_to_tensor(input_ids, dtype=tf.int32), 0),
            speaker_ids=tf.convert_to_tensor([0], dtype=tf.int32),
            speed_ratios=tf.convert_to_tensor([1.0], dtype=tf.float32),
            f0_ratios=tf.convert_to_tensor([1.0], dtype=tf.float32),
            energy_ratios=tf.convert_to_tensor([1.0], dtype=tf.float32),
        )
    else:
        raise ValueError("Only TACOTRON, FASTSPEECH2 are supported on text2mel_name")

    # vocoder part
    if vocoder_name == "MB-MELGAN":
        # tacotron-2 generate noise in the end symtematic, let remove it :v.
        if text2mel_name == "TACOTRON":
            remove_end = 1024
        else:
            remove_end = 1
        audio = vocoder_model.inference(mel_outputs)[0, :-remove_end, 0]
    else:
        raise ValueError("Only MB_MELGAN are supported on vocoder_name")

    if text2mel_name == "TACOTRON":
        return mel_outputs.numpy(), alignment_history.numpy(), audio.numpy()
    else:
        return mel_outputs.numpy(), audio.numpy()


def run_fun_in_new_thread(f, args=()):
    threading.Thread(target=f, args=args).start()


task_list = []


def loop():
    time.sleep(5)
    while True:
        if task_list:
            audios = task_list.pop()
            # run_fun_in_new_thread(f=lambda *audios: [sd.play(audio, 24000, blocking=True) for audio in audios],
            #                       args=audios)
            _ = [sd.play(audio, 24000, blocking=True) for audio in audios]
        time.sleep(1)


def spl_str(txt: str):
    out = re.split('\n|,|，', txt)
    return out


@app.route("/tts")
def tts(request):
    input_text = request.args['message'][0]
    # 每行读一次，不然不行
    audios = []
    line_short_msgs = spl_str(str(input_text))
    for short_msg in line_short_msgs:
        mels, audio = do_synthesis(short_msg, fastspeech2, mb_melgan, "FASTSPEECH2", "MB-MELGAN")
        audios.append(audio)

    task_list.append(audios)
    # sf.write('test.wav',wav,24000)
    # sd.stop()
    # run_fun_in_new_thread(f=lambda *audios: [sd.play(audio, 24000, blocking=True) for audio in audios], args=audios)

    return response.json(
        {
            "parsed": True,
            "args": request.args,
            "url": request.url,
            "query_string": request.query_string,
        }
    )


if __name__ == '__main__':
    run_fun_in_new_thread(loop, args=())
    app.run(host="0.0.0.0", port=8000, debug=True)
